import bpy
import math
import os
import bpy_extras.object_utils

def create_cube(location, size=2):
    bpy.ops.mesh.primitive_cube_add(size=size, location=location)
    cube = bpy.context.object
    mat = bpy.data.materials.new(name="WhiteMaterial")
    mat.diffuse_color = (1, 1, 1, 1)  # 纯白色
    cube.data.materials.append(mat)
    return cube

def create_black_circle(location, radius=0.1):
    bpy.ops.mesh.primitive_circle_add(radius=radius, location=location, fill_type='NGON')
    circle = bpy.context.object
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.fill()  # 填充圆形
    bpy.ops.object.mode_set(mode='OBJECT')
    mat = bpy.data.materials.new(name="BlackMaterial")
    mat.diffuse_color = (0, 0, 0, 1)  # 纯黑色
    circle.data.materials.append(mat)
    return circle

def create_camera(location, rotation, lens=50, sensor_width=36, sensor_height=24):
    bpy.ops.object.camera_add(location=location, rotation=rotation)
    camera = bpy.context.object
    camera.data.lens = lens
    camera.data.sensor_width = sensor_width
    camera.data.sensor_height = sensor_height
    return camera

def get_object_screen_position(obj, camera):
    scene = bpy.context.scene
    scene.camera = camera
    bpy.context.view_layer.update()
    depsgraph = bpy.context.evaluated_depsgraph_get()
    obj_eval = obj.evaluated_get(depsgraph)
    co = obj_eval.matrix_world.translation
    co_2d = bpy_extras.object_utils.world_to_camera_view(scene, camera, co)
    render_scale = scene.render.resolution_percentage / 100
    render_size = (
        int(scene.render.resolution_x * render_scale),
        int(scene.render.resolution_y * render_scale),
    )
    return (round(co_2d.x * render_size[0]), round((1.0 - co_2d.y) * render_size[1]))

def get_camera_parameters(camera):
    # 获取相机外参（位置和旋转矩阵）
    location = camera.location
    rotation = camera.rotation_euler.to_matrix().to_4x4()

    # 获取相机内参（焦距和传感器尺寸）
    focal_length = camera.data.lens
    sensor_width = camera.data.sensor_width
    sensor_height = camera.data.sensor_height

    # 计算相机矩阵K
    scene = bpy.context.scene
    render = scene.render
    resolution_x_in_px = render.resolution_x
    resolution_y_in_px = render.resolution_y
    scale = render.resolution_percentage / 100
    resolution_x_in_px = int(scale * resolution_x_in_px)
    resolution_y_in_px = int(scale * resolution_y_in_px)

    pixel_aspect_ratio = render.pixel_aspect_x / render.pixel_aspect_y
    if (camera.data.sensor_fit == 'VERTICAL'):
        s_u = resolution_x_in_px / sensor_width / pixel_aspect_ratio
        s_v = resolution_y_in_px / sensor_height
    else:  # 'HORIZONTAL' and 'AUTO'
        s_u = resolution_x_in_px / sensor_width
        s_v = resolution_y_in_px * pixel_aspect_ratio / sensor_height

    u_0 = resolution_x_in_px / 2
    v_0 = resolution_y_in_px / 2

    K = [[s_u * focal_length, 0, u_0],
         [0, s_v * focal_length, v_0],
         [0, 0, 1]]

    # 获取畸变系数（假设零畸变）
    distortion_coeffs = [0, 0, 0, 0, 0]

    return location, rotation, K, distortion_coeffs

def setup_scene():
    # 创建两个正方体
    cube1 = create_cube((0, 0, 1))
    cube2 = create_cube((4, 0, 1))

    # 获取正方体顶面的坐标，并在其上放置黑色圆形区域
    circle1 = create_black_circle((0, 0, 2.01))  # 在第一个正方体顶面放置黑色圆形，略高于表面
    circle2 = create_black_circle((4, 0, 2.01))  # 在第二个正方体顶面放置黑色圆形，略高于表面

    # 创建相机，模拟成人眼睛位置，从上方拍摄
    camera_left = create_camera((2 - 0.03, 0, 15), (math.radians(0), 0, 0), lens=35, sensor_width=32,
                                sensor_height=18)  # 调整相机位置和角度
    camera_right = create_camera((2 + 0.03, 0, 15), (math.radians(0), 0, 0), lens=50, sensor_width=36,
                                 sensor_height=24)  # 调整相机位置和角度

    return cube1, cube2, circle1, circle2, camera_left, camera_right

def render_scene(camera, output_path):
    bpy.context.scene.camera = camera
    bpy.context.scene.render.filepath = output_path
    bpy.ops.render.render(write_still=True)

if __name__ == "__main__":
    import bpy_extras.object_utils

    # 设置场景
    cube1, cube2, circle1, circle2, camera_left, camera_right = setup_scene()

    # 指定保存目录
    output_folder = "E:/Blender/output/output_blender_test"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 渲染左相机视角
    render_scene(camera_left, os.path.join(output_folder, "left_camera.png"))

    # 渲染右相机视角并计算视差
    pos_left_circle1 = get_object_screen_position(circle1, camera_left)
    pos_right_circle1 = get_object_screen_position(circle1, camera_right)
    disparity_circle1 = abs(pos_left_circle1[0] - pos_right_circle1[0])

    pos_left_circle2 = get_object_screen_position(circle2, camera_left)
    pos_right_circle2 = get_object_screen_position(circle2, camera_right)
    disparity_circle2 = abs(pos_left_circle2[0] - pos_right_circle2[0])

    # 渲染右相机视角
    render_scene(camera_right, os.path.join(output_folder, "right_camera.png"))

    # 获取左相机参数
    bpy.context.scene.camera = camera_left
    bpy.context.view_layer.update()
    left_camera_params = get_camera_parameters(camera_left)

    # 获取右相机参数
    bpy.context.scene.camera = camera_right
    bpy.context.view_layer.update()
    right_camera_params = get_camera_parameters(camera_right)

    # 生成包含视差结果和相机参数的Python文件
    try:
        with open(os.path.join(output_folder, "disparity_result.py"), 'w') as f:
            f.write(f"# 视差结果\n")
            f.write(f"disparity_circle1 = {disparity_circle1}\n")
            f.write(f"disparity_circle2 = {disparity_circle2}\n")
            f.write("print(f'视差 (circle 1): {disparity_circle1} 像素')\n")
            f.write("print(f'视差 (circle 2): {disparity_circle2} 像素')\n")
            f.write("\n# 左相机参数\n")
            f.write(f"left_camera_location = {left_camera_params[0]}\n")
            f.write(f"left_camera_rotation = {left_camera_params[1]}\n")
            f.write(f"left_camera_intrinsics = {left_camera_params[2]}\n")
            f.write(f"left_camera_distortion_coeffs = {left_camera_params[3]}\n")
            f.write("\n# 右相机参数\n")
            f.write(f"right_camera_location = {right_camera_params[0]}\n")
            f.write(f"right_camera_rotation = {right_camera_params[1]}\n")
            f.write(f"right_camera_intrinsics = {right_camera_params[2]}\n")
            f.write(f"right_camera_distortion_coeffs = {right_camera_params[3]}\n")
        print(f"Python file with disparity results and camera parameters created at {os.path.join(output_folder, 'disparity_result.py')}")
    except Exception as e:
        print(f"Failed to create Python file: {e}")

    print(f"Rendering complete. Images and disparity result saved to {output_folder} directory.")
